#include "benchmark_base.h"
#include <yams/api/content_store.h>
#include <yams/search/search_executor.h>
#include <yams/search/hybrid_search_engine.h>
#include "../common/test_data_generator.h"
#include "../common/fixture_manager.h"
#include <filesystem>
#include <random>

using namespace yams;
using namespace yams::benchmark;

class SearchBenchmark : public BenchmarkBase {
public:
    SearchBenchmark(const std::string& name, size_t corpusSize) 
        : BenchmarkBase("Search_" + name), corpusSize_(corpusSize) {
        setUp();
    }
    
    ~SearchBenchmark() {
        tearDown();
    }

protected:
    void setUp() {
        // Create temp directory
        tempDir_ = std::filesystem::temp_directory_path() / "yams_bench_search";
        std::filesystem::create_directories(tempDir_);
        
        // Initialize components
        contentStore_ = std::make_unique<api::ContentStore>(tempDir_.string());
        
        // Build corpus
        buildCorpus();
        
        // Generate search queries
        generateQueries();
    }
    
    void tearDown() {
        std::filesystem::remove_all(tempDir_);
    }
    
    void buildCorpus() {
        std::cout << "Building corpus of " << corpusSize_ << " documents..." << std::endl;
        
        test::TestDataGenerator generator;
        
        // Add variety of document sizes
        for (size_t i = 0; i < corpusSize_; ++i) {
            size_t docSize = 1024 + (i % 10) * 1024; // 1KB to 10KB
            auto content = generator.generateTextDocument(docSize);
            
            // Add some searchable patterns
            if (i % 10 == 0) content += " benchmark test document";
            if (i % 5 == 0) content += " performance evaluation";
            if (i % 3 == 0) content += " search query optimization";
            
            contentStore_->addContent(
                std::vector<std::byte>(content.begin(), content.end()),
                "doc_" + std::to_string(i) + ".txt"
            );
        }
    }
    
    void generateQueries() {
        queries_ = {
            "benchmark",
            "performance",
            "search",
            "optimization",
            "test document",
            "evaluation metrics",
            "query performance",
            "benchmark test",
            "search optimization",
            "performance evaluation"
        };
    }
    
    void collectCustomMetrics(std::map<std::string, double>& metrics) override {
        metrics["corpus_size"] = corpusSize_;
        metrics["avg_results"] = avgResultCount_;
        metrics["avg_score"] = avgScore_;
    }
    
    std::filesystem::path tempDir_;
    std::unique_ptr<api::ContentStore> contentStore_;
    size_t corpusSize_;
    std::vector<std::string> queries_;
    
    double avgResultCount_ = 0;
    double avgScore_ = 0;
};

// Exact match search in small corpus (1K docs)
BENCHMARK_F(SearchBenchmark, ExactMatch_1K) {
    SearchBenchmark("ExactMatch_1K", 1000) {}
    
    size_t runIteration() override {
        size_t totalResults = 0;
        double totalScore = 0;
        
        for (const auto& query : queries_) {
            auto results = contentStore_->search(query);
            if (results.has_value()) {
                totalResults++;
                totalScore += results->score;
            }
        }
        
        avgResultCount_ = totalResults / static_cast<double>(queries_.size());
        avgScore_ = totalScore / queries_.size();
        
        return queries_.size(); // Number of searches performed
    }
};

// Exact match search in medium corpus (10K docs)
BENCHMARK_F(SearchBenchmark, ExactMatch_10K) {
    SearchBenchmark("ExactMatch_10K", 10000) {}
    
    size_t runIteration() override {
        size_t totalResults = 0;
        
        for (const auto& query : queries_) {
            auto results = contentStore_->search(query);
            if (results.has_value()) {
                totalResults++;
            }
        }
        
        avgResultCount_ = totalResults / static_cast<double>(queries_.size());
        return queries_.size();
    }
};

// Fuzzy search benchmark
class FuzzySearchBenchmark : public SearchBenchmark {
public:
    FuzzySearchBenchmark() : SearchBenchmark("FuzzySearch", 5000) {}
    
protected:
    size_t runIteration() override {
        size_t totalResults = 0;
        
        // Test with typos and variations
        std::vector<std::string> fuzzyQueries = {
            "benchmrk",      // Typo
            "perfrmance",    // Missing letter
            "serach",        // Transposition
            "optimizaton",   // Missing letter
            "documnet",      // Typo
            "evalutation"    // Extra letter
        };
        
        for (const auto& query : fuzzyQueries) {
            search::SearchOptions options;
            options.fuzzy = true;
            options.similarity = 0.8;
            
            auto results = contentStore_->searchWithOptions(query, options);
            totalResults += results.size();
        }
        
        avgResultCount_ = totalResults / static_cast<double>(fuzzyQueries.size());
        return fuzzyQueries.size();
    }
};

// Metadata filtering benchmark
class MetadataFilterBenchmark : public SearchBenchmark {
public:
    MetadataFilterBenchmark() : SearchBenchmark("MetadataFilter", 5000) {}
    
protected:
    void setUp() {
        SearchBenchmark::setUp();
        
        // Add metadata to documents
        test::TestDataGenerator generator;
        for (size_t i = 0; i < corpusSize_; ++i) {
            auto metadata = generator.generateMetadata(5);
            // Would need metadata repository access to set metadata
            // For benchmark purposes, we're measuring search with filters
        }
    }
    
    size_t runIteration() override {
        size_t totalResults = 0;
        
        // Search with metadata filters
        std::vector<std::string> filters = {
            "status:completed",
            "priority:high",
            "type:document",
            "author:test",
            "tag:benchmark"
        };
        
        for (const auto& filter : filters) {
            search::SearchOptions options;
            options.metadataFilter = filter;
            
            auto results = contentStore_->searchWithOptions("*", options);
            totalResults += results.size();
        }
        
        avgResultCount_ = totalResults / static_cast<double>(filters.size());
        return filters.size();
    }
};

// Combined text and metadata search
class CombinedSearchBenchmark : public SearchBenchmark {
public:
    CombinedSearchBenchmark() : SearchBenchmark("Combined", 5000) {}
    
protected:
    size_t runIteration() override {
        size_t totalResults = 0;
        
        struct SearchQuery {
            std::string text;
            std::string metadataFilter;
            bool fuzzy;
        };
        
        std::vector<SearchQuery> queries = {
            {"benchmark", "status:active", false},
            {"performance", "priority:high", true},
            {"test", "type:document", false},
            {"optimization", "tag:important", true},
            {"search", "author:system", false}
        };
        
        for (const auto& query : queries) {
            search::SearchOptions options;
            options.fuzzy = query.fuzzy;
            options.metadataFilter = query.metadataFilter;
            
            auto results = contentStore_->searchWithOptions(query.text, options);
            totalResults += results.size();
        }
        
        avgResultCount_ = totalResults / static_cast<double>(queries.size());
        return queries.size();
    }
};

// Pagination benchmark
class PaginationBenchmark : public SearchBenchmark {
public:
    PaginationBenchmark() : SearchBenchmark("Pagination", 10000) {}
    
protected:
    size_t runIteration() override {
        const size_t pageSize = 20;
        size_t totalPages = 0;
        
        // Search and paginate through results
        search::SearchOptions options;
        options.limit = pageSize;
        
        for (const auto& query : queries_) {
            size_t offset = 0;
            size_t pageCount = 0;
            
            while (pageCount < 5) { // Limit to 5 pages per query
                options.offset = offset;
                auto results = contentStore_->searchWithOptions(query, options);
                
                if (results.empty()) break;
                
                pageCount++;
                offset += pageSize;
            }
            
            totalPages += pageCount;
        }
        
        avgResultCount_ = totalPages;
        return totalPages; // Number of pages retrieved
    }
};

// Concurrent search benchmark
class ConcurrentSearchBenchmark : public SearchBenchmark {
public:
    ConcurrentSearchBenchmark() : SearchBenchmark("Concurrent", 5000) {}
    
protected:
    size_t runIteration() override {
        const size_t numThreads = 4;
        std::atomic<size_t> totalSearches{0};
        std::atomic<size_t> totalResults{0};
        
        std::vector<std::thread> threads;
        for (size_t t = 0; t < numThreads; ++t) {
            threads.emplace_back([this, t, &totalSearches, &totalResults]() {
                std::mt19937 rng(t);
                std::uniform_int_distribution<size_t> dist(0, queries_.size() - 1);
                
                for (size_t i = 0; i < 25; ++i) {
                    size_t queryIdx = dist(rng);
                    auto results = contentStore_->search(queries_[queryIdx]);
                    
                    totalSearches++;
                    if (results.has_value()) {
                        totalResults++;
                    }
                }
            });
        }
        
        for (auto& thread : threads) {
            thread.join();
        }
        
        avgResultCount_ = totalResults.load() / static_cast<double>(totalSearches.load());
        return totalSearches.load();
    }
};

// Main benchmark runner
int main(int argc, char** argv) {
    BenchmarkBase::Config config;
    config.verbose = true;
    config.benchmark_iterations = 5;
    config.warmup_iterations = 2;
    
    // Parse arguments
    for (int i = 1; i < argc; ++i) {
        std::string arg = argv[i];
        if (arg == "--quiet") {
            config.verbose = false;
        } else if (arg == "--iterations" && i + 1 < argc) {
            config.benchmark_iterations = std::stoi(argv[++i]);
        } else if (arg == "--output" && i + 1 < argc) {
            config.output_file = argv[++i];
        }
    }
    
    // Run benchmarks
    std::cout << "YAMS Search Performance Benchmarks\n";
    std::cout << "===================================\n\n";
    
    std::vector<std::unique_ptr<BenchmarkBase>> benchmarks;
    
    // Create and configure benchmarks
    auto addBenchmark = [&](BenchmarkBase* bench) {
        bench->config = config;
        benchmarks.emplace_back(bench);
    };
    
    addBenchmark(new ExactMatch_1KBenchmark());
    addBenchmark(new ExactMatch_10KBenchmark());
    addBenchmark(new FuzzySearchBenchmark());
    addBenchmark(new MetadataFilterBenchmark());
    addBenchmark(new CombinedSearchBenchmark());
    addBenchmark(new PaginationBenchmark());
    addBenchmark(new ConcurrentSearchBenchmark());
    
    // Track results
    test::BenchmarkTracker tracker("search_benchmarks.json");
    
    for (auto& benchmark : benchmarks) {
        std::cout << "\n";
        auto result = benchmark->run();
        
        // Record for tracking
        test::BenchmarkTracker::BenchmarkResult trackerResult;
        trackerResult.name = result.name;
        trackerResult.value = result.duration_ms;
        trackerResult.unit = "ms";
        trackerResult.timestamp = std::chrono::system_clock::now();
        trackerResult.metrics = result.custom_metrics;
        
        tracker.recordResult(trackerResult);
    }
    
    // Generate reports
    tracker.generateReport("search_benchmark_report.json");
    tracker.generateMarkdownReport("search_benchmark_report.md");
    
    std::cout << "\n===================================\n";
    std::cout << "Benchmark complete. Reports generated.\n";
    
    return 0;
}